data <- read.csv("data.csv") 
data$crisis <- 1*(data$growthc_t1 < -1)
data$mediator <- data$nelda12
data$s18f1_lag <- 1*(data$s18f1_lag > 0)
data <- data[ ,c("mirt", "country", "year", "mediator", "crisis", "lambda", "agedem_t1", "nelda19", "legelec", "s18f1_lag", "gdp.gle_t1", "nelda45", "nelda53")]

# Remove NA's (not allowed by mediate)
data <- na.omit(data)

Model <- mirt ~ crisis + lambda + s18f1_lag + gdp.gle_t1 + agedem_t1 + nelda19 + legelec + (1|country)

#INSTRUMENT STRENGTH
hatfit0 <- lmer(Model, data = data)
hatfit  <- update(hatfit0, . ~ . + nelda45 + nelda53, na.action = na.exclude)

# F-statistic for first stage = Chisq/2
anova(hatfit, hatfit0)$Chisq[2]/2

#COMPUTE PREDICTED MANIPULATION VALUES
data$mirt.hat <- predict(hatfit)

# Equation 3.1 (treatement effect on the mediator)
model.m <- glmer(update(Model, mediator ~ . + mirt.hat), family = binomial(probit), data = data, glmerControl(optimizer = "bobyqa", optCtrl = list(maxfun = 1e+6)))

# Equation 3.2 (mediator effect on the outcome)
model.y <- update(hatfit0, . ~ . + mediator)
  
m.out <- mediate(model.m, model.y, treat="crisis", mediator="mediator", sim=1000)

print(summary(m.out))

pdf("../graphs/mediation.pdf", width=6, height=3)
par(mar=c(3, 14, 1, 1))
plot(rev(c(m.out$d0, m.out$z0, m.out$tau.coef)), 1:3, xlim = c(-.25, .15), ylim=c(0.5,3.5), pch = 16, cex = 1.5, yaxt = 'n', ylab = "", xlab = "", cex.lab = 1.5)
ci <- t(cbind(m.out$d0.ci, m.out$z0.ci, m.out$tau.ci))
segments(rev(ci[,1]), 1:3, rev(ci[,2]), 1:3, lwd = 3)
abline(v = 0, lty = 3, col = "grey", lwd = 3)
axis(2, at = 1:3, rev(c("Average mediation effect (ACME)", "Average direct effect (ADE)", "Total effect")), las = 1)
dev.off()

# Remove country-effects for sensitivity analysis (medsens does not support lmer objects)

model.m <- glm(update(Model, mediator ~ . + mirt.hat - (1|country)), family = binomial(probit), data = data)

# Equation 3.2 (mediator effect on the outcome)
model.y <- lm(update(Model, . ~ . + mediator - (1|country)), data = data)
m.out <- mediate(model.m, model.y, treat="crisis", mediator="mediator", sim=1000)
sens <- medsens(m.out, rho.by = 0.1, sims=500)

cat("product threshold:", sqrt(sens$R2star.d.thresh), "\n")

pdf("../graphs/sensitivity.pdf", width=10, height=5)
par(mfrow=c(1,2))
plot(sens, sens.par = "rho")
plot(sens, sens.par = "R2")
dev.off()

